# %% Import relevant modules
from matplotlib import pyplot
from scipy.stats import zscore
import pandas as pd
import numpy as np
import warnings
import os
from sklearn.decomposition import PCA
from sklearn.impute import KNNImputer
warnings.filterwarnings("ignore")
# %% Loading data
path = "./EquityData"
files = os.listdir(path)
index = '000852.XSHG'
# file extensions
market_info_extension = "_Prices.csv"
index_info_extension = "_index_information.csv"
industry_info_extension = "_Industry.csv"
fundamental_info_extension = "_fundamentals.csv"
pricing_factors_info_extension = "_total_factors.csv"
valuation_info_extension = "valuation_information.csv"

# filter file names
list_market_info = [elements for elements in files if market_info_extension in elements]
list_index_info = [elements for elements in files if index_info_extension in elements]
list_industry_info = [elements for elements in files if industry_info_extension in elements]
list_fundamental_info = [elements for elements in files if fundamental_info_extension in elements]
list_pricing_factors_info = [elements for elements in files if pricing_factors_info_extension in elements]
list_valuation_info = [elements for elements in files if valuation_info_extension in elements]

# load files
market_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_market_info]
market_info_df = pd.concat(market_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)

index_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_index_info]
index_info_df = pd.concat(index_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)

industry_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_industry_info]
industry_info_df = pd.concat(industry_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)

pricing_factors_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_pricing_factors_info]
pricing_factors_info_df = pd.concat(pricing_factors_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)
pricing_factor_names = pricing_factors_info_df.columns[2:-2]

# marge data frames
meta_data = pd.merge(
    market_info_df,
    index_info_df,
    on=["time", "code"],
    how="inner"
)
meta_data = pd.merge(
    meta_data,
    industry_info_df,
    on=["time", "code"],
    how="inner"
)
meta_data = pd.merge(
    meta_data,
    pricing_factors_info_df,
    on=["time", "code"],
    how="inner"
)

meta_data = meta_data.sort_values(['time', 'code']).reset_index(drop=True)
meta_data = meta_data.sort_values(['time', 'code']).reset_index(drop=True)
del files, fundamental_info_extension, index_info_extension, industry_info_extension\
    , market_info_extension, pricing_factors_info_extension, valuation_info_extension,list_market_info, list_valuation_info, list_fundamental_info\
    ,list_index_info, list_industry_info, list_pricing_factors_info

# 描述性统计
print(meta_data[meta_data.columns[:20]].info())
meta_data.dropna(axis=1, inplace=True,  thresh=len(meta_data)*0.8) # thresh
meta_data.drop(columns=['index', 'industry'], inplace=True)
meta_data.groupby(['code']).shift(0).fillna(method='bfill', inplace=True)
# 缺失值处理
imputer = KNNImputer(n_neighbors=10)
meta_data.replace([np.inf, -np.inf], np.nan, inplace=True)
meta_data_imputer = pd.DataFrame(imputer.fit_transform(meta_data[meta_data.columns[2:]]), columns=meta_data.columns[2:], index=meta_data.index)
# meta_data_imputer.isnull().mean().sum()
# PCA
pca = PCA(n_components=25)
pca_meta_data = pd.DataFrame(pca.fit_transform(meta_data_imputer))
pca_meta_data = pd.concat([meta_data[meta_data.columns[:2]], pca_meta_data], axis=1)
pca_meta_data.to_csv('./' + index + '_meta_data.csv')